import torch.nn as nn
import torch
import numpy as np
import pandas as pd
from django.http import HttpResponse, JsonResponse
from django.views.decorators.csrf import csrf_exempt

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

import pandas as pd
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
import os
import tushare as ts

ts.set_token('084256242bd0568c9ea0786ff7019ea73e1d83efa7d51a3454f99d67')
pro = ts.pro_api()


def get_data(code, start, end):
    # end1 = str(int(end) + 1)
    df = pro.daily(ts_code=code, autype='qfq', start_date=start, end_date=end)
    # print(df)
    df.index = pd.to_datetime(df.trade_date)
    # 设置把日期作为索引
    # df['ma'] = 0.0  # Backtrader需要用到
    # df['openinterest'] = 0.0  # Backtrader需要用到
    # 定义两个新的列ma和openinterest
    df = df[['open', 'high', 'low', 'close', 'vol']]
    # 重新设置df取值，并返回df
    return df


def acquire_code(inp_code, inp_start, inp_end):  # 只下载一只股票数据，且只用CSV保存   未来可以有自己的数据库
    df = get_data(inp_code, inp_start, inp_end)
    # print(df.info())
    # 输出统计各列的数据量
    # print("—" * 30)
    # 分割线
    # print(df.describe())
    # 输出常用统计参数
    df.sort_index(inplace=True)
    # 把股票数据按照时间正序排列
    path = os.path.join(os.path.join(os.getcwd(),
                                     "stock\dataset"), inp_code + ".csv")
    # os.path地址拼接，''数据地址''为文件保存路径
    # path = os.path.join(os.path.join(os.getcwd(),"数据地址"),inp_code+"_30M.csv")
    df.to_csv(path)



import csv




@csrf_exempt
def lstm(request):
    code = request.GET.get("code")
    code2 = request.GET.get("code2")
    code3 = request.GET.get("code3")
    acquire_code(code, code2, code3)

    filepath = './stock/dataset/' + code + '.csv'
    data = pd.read_csv(filepath)
    data = data.sort_values('trade_date')
    print(data.head())
    print(data.shape)

    # 1.特征工程
    # 选取Close作为特征
    price = data[['close']]
    print(price.info())

    from sklearn.preprocessing import MinMaxScaler

    # 进行不同的数据缩放，将数据缩放到-1和1之间
    scaler = MinMaxScaler(feature_range=(-1, 1))
    price['close'] = scaler.fit_transform(price['close'].values.reshape(-1, 1))
    print(price['close'].shape)
    print(price['close'])

    # 2.数据集制作
    # 今天的收盘价预测明天的收盘价
    # lookback表示观察的跨度
    def split_data(stock, lookback):
        data_raw = stock.to_numpy()
        data = []
        # print(data)

        # you can free play（seq_length）
        for index in range(len(data_raw) - lookback):
            data.append(data_raw[index: index + lookback])

        data = np.array(data)
        test_set_size = int(np.round(0.2 * data.shape[0]))
        train_set_size = data.shape[0] - (test_set_size)

        x_train = data[:train_set_size, :-1, :]
        y_train = data[:train_set_size, -1, :]

        x_test = data[train_set_size:, :-1]
        y_test = data[train_set_size:, -1, :]

        return [x_train, y_train, x_test, y_test]

    lookback = 20
    x_train, y_train, x_test, y_test = split_data(price, lookback)

    # 注意：pytorch的nn.LSTM input shape=(seq_length, batch_size, input_size)
    # 3.模型构建 —— LSTM

    import torch
    import torch.nn as nn

    x_train = torch.from_numpy(x_train).type(torch.Tensor)
    x_test = torch.from_numpy(x_test).type(torch.Tensor)
    y_train_lstm = torch.from_numpy(y_train).type(torch.Tensor)
    y_test_lstm = torch.from_numpy(y_test).type(torch.Tensor)
    y_train_gru = torch.from_numpy(y_train).type(torch.Tensor)
    y_test_gru = torch.from_numpy(y_test).type(torch.Tensor)
    # 输入的维度为1，只有Close收盘价
    input_dim = 1
    # 隐藏层特征的维度
    hidden_dim = 32
    # 循环的layers
    num_layers = 2
    # 预测后一天的收盘价
    output_dim = 1
    num_epochs = 100

    class LSTM(nn.Module):
        def __init__(self, input_dim, hidden_dim, num_layers, output_dim):
            super(LSTM, self).__init__()
            self.hidden_dim = hidden_dim
            self.num_layers = num_layers

            self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True)
            self.fc = nn.Linear(hidden_dim, output_dim)

        def forward(self, x):
            h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).requires_grad_()
            c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).requires_grad_()
            out, (hn, cn) = self.lstm(x, (h0.detach(), c0.detach()))
            out = self.fc(out[:, -1, :])
            return out

    model = LSTM(input_dim=input_dim, hidden_dim=hidden_dim, output_dim=output_dim, num_layers=num_layers)
    criterion = torch.nn.MSELoss()
    optimiser = torch.optim.Adam(model.parameters(), lr=0.01)

    # 4.模型训练
    import time

    hist = np.zeros(num_epochs)
    start_time = time.time()
    lstm = []

    for t in range(num_epochs):
        y_train_pred = model(x_train)

        loss = criterion(y_train_pred, y_train_lstm)
        print("Epoch ", t, "MSE: ", loss.item())
        hist[t] = loss.item()

        optimiser.zero_grad()
        loss.backward()
        optimiser.step()

    training_time = time.time() - start_time
    print("Training time: {}".format(training_time))

    # 5.模型验证
    # print(x_test[-1])
    import math, time
    from sklearn.metrics import mean_squared_error

    # make predictions
    y_test_pred = model(x_test)

    # invert predictions
    y_train_pred = scaler.inverse_transform(y_train_pred.detach().numpy())
    y_train = scaler.inverse_transform(y_train_lstm.detach().numpy())
    y_test_pred = scaler.inverse_transform(y_test_pred.detach().numpy())
    y_test = scaler.inverse_transform(y_test_lstm.detach().numpy())

    # calculate root mean squared error
    trainScore = math.sqrt(mean_squared_error(y_train[:, 0], y_train_pred[:, 0]))
    print('Train Score: %.2f RMSE' % (trainScore))
    testScore = math.sqrt(mean_squared_error(y_test[:, 0], y_test_pred[:, 0]))
    print('Test Score: %.2f RMSE' % (testScore))
    lstm.append(trainScore)
    lstm.append(testScore)
    lstm.append(training_time)

    # In[40]:

    # shift train predictions for plotting
    trainPredictPlot = np.empty_like(price)
    trainPredictPlot[:, :] = np.nan
    trainPredictPlot[lookback:len(y_train_pred) + lookback, :] = y_train_pred

    # shift test predictions for plotting
    testPredictPlot = np.empty_like(price)
    testPredictPlot[:, :] = np.nan
    testPredictPlot[len(y_train_pred) + lookback:len(price), :] = y_test_pred

    original = scaler.inverse_transform(price['close'].values.reshape(-1, 1))

    predictions = np.append(trainPredictPlot, testPredictPlot, axis=1)
    predictions = np.append(predictions, original, axis=1)
    result = pd.DataFrame(predictions)
    result[2].iloc[-1] = 'NaN'
    tomorrow_price = str(result[1].iloc[-1])
    print("明天的股票价格为:" + tomorrow_price)

    import datetime
    c = 0
    data2 = data['trade_date'].copy()
    # print(data2)
    for i in data['trade_date']:
        str_time = datetime.datetime.strptime(i, "%Y-%m-%d")
        data2.iloc[c] = (str_time + datetime.timedelta(days=1)).strftime("%Y-%m-%d")
        c += 1
    # print(data2)

    # 折线图
    import plotly.express as px
    import plotly.graph_objects as go
    import plotly.offline as offline

    fig = go.Figure()
    fig.add_trace(go.Scatter(go.Scatter(x=result.index, y=result[0],
                                        mode='lines',
                                        name='Train prediction')))
    fig.add_trace(go.Scatter(x=result.index, y=result[1],
                             mode='lines',
                             name='Test prediction'))
    fig.add_trace(go.Scatter(go.Scatter(x=result.index, y=result[2],
                                        mode='lines',
                                        name='Actual Value')))

    fig.update_layout(
        xaxis=dict(
            showline=True,
            showgrid=False,
            showticklabels=False,
            linecolor='white',
            linewidth=2,
            tickvals=result.index[::1],
            ticktext=data2.loc[::1],
        ),
        yaxis=dict(
            title_text='Close (USD)',
            titlefont=dict(
                family='Rockwell',
                size=12,
                color='white',
            ),
            showline=True,
            showgrid=True,
            showticklabels=True,
            linecolor='white',
            linewidth=2,
            ticks='outside',
            tickfont=dict(
                family='Rockwell',
                size=12,
                color='white',
            ),
        ),
        showlegend=True,
        template='plotly_dark'

    )

    annotations = []
    annotations.append(dict(xref='paper', yref='paper', x=0.0, y=1.05,
                            xanchor='left', yanchor='bottom',
                            text='Results (LSTM)'+"股票代码：" + code,
                            font=dict(family='Rockwell',
                                      size=26,
                                      color='white'),
                            showarrow=False))
    fig.update_layout(annotations=annotations)
    offline.plot(fig, filename='stock/stock_prediction_html/stock_prediction.html', auto_open=False)
    # fig.show()
    return JsonResponse({'code': 200, 'msg': 'success', 'data': tomorrow_price})